file_path = r'dataset2.txt'
support_num = 3  # 支持度设置
confidence = 0.7  # 置信度设置


def read_data(path: str)->list:
    """
    从文件中读取数据
    :param path: 文件路径
    :return: 返回一个列表，每个列表就是一行
    """
    with open(path, 'r', encoding='utf8') as f:
        return f.readlines()[1::]  # 去掉第一行


def deal_data(data_list: list)->list:
    """
    处理每行数据，返回一个列表，列表中每个值是该行的一个set
    :param data_list:
    :return:
    """
    data = list()
    import re
    for line in data_list:
        split_list = re.split(r',?\s', line)[1::]
        item_set = set()
        for i in split_list:
            if i:  # 上面split分割的数据最后会有一个''，因为原始数据最后有\n，这里需要去除一下空字符
                item_set.add(i)
        data.append(item_set)
    return data


def get_all_item(data: list)->list:
    """
    返回事物数据中的一项集set列表，格式为[{item1}, {item2}...]
    :param data: 事物数据
    :return: 一项集set列表，格式为[{item1}, {item2}...]
    """
    item_set_list = list()
    for s in data:
        for i in s:
            if {i} not in item_set_list:
                item_set_list.append({i})
    return item_set_list


def get_support(items: set, data: list)->int:
    """
    传入一个项集和事物数据，返回项集在事物数据中的支持度计数
    :param items: 项集set
    :param data: 事物数据
    :return: 支持度计数
    """
    support_count = 0
    for t in data:
        if items & t == items:  # 项集set和事物数据中一行取交集，如果交集等于项集set本身，说明项集中的所有元素在事物数据的该行中都出现
            support_count += 1
    return support_count


def get_frequent_items(candi_items: list, data: list, min_support: int)->list:
    """
    获取频繁项集set的list
    :param candi_items: 候选项集list
    :param data: 事物数据
    :param min_support: 最小支持度
    :return: 频繁项集set的list
    """
    res = list()
    for i in candi_items:
        if get_support(i, data) >= min_support:
            res.append(i)
    return res


def get_candi_items(items: list, item_num: int)->list:
    """
    获取项数+1的候选项集
    :param items: 当前项
    :param item_num: 需要生成的候选项集项数
    :return: 项数+1的候选项集
    """
    # 生成所有候选项集
    all_candi_items = list()
    for i in range(len(items) ** 2):
        items_set = set()
        for j in range(len(items)):
            if i & (1 << j):
                for k in items[j]:
                    items_set.add(k)
        if len(items_set) == item_num:
            all_candi_items.append(items_set)

    # 剪枝
    final_candi_items = list()
    for i in items:
        for j in all_candi_items:
            if i & j:
                if j not in final_candi_items:
                    final_candi_items.append(j)
    return final_candi_items


def get_all_freq_items(data: list, support: int)->list:
    """
    返回数据中的频繁项集的列表，列表元素0为1项集，元素1为2项集...
    :param data: 事物数据
    :param support: 支持度
    :return: 列表元素0为1项集，元素1为2项集...
    """
    freq_items_list = list()
    candi_items = get_all_item(data)  # 获取候选一项集
    item_count = 1
    while True:
        freq_items = get_frequent_items(candi_items, data, support)  # 获取频繁项集
        if freq_items:
            # print('频繁{}项集：{}'.format(item_count, freq_items))
            freq_items_list.append(freq_items)
            item_count += 1
            candi_items = get_candi_items(freq_items, item_count)  # 在频繁项集基础上获取项数+1的候选项集
            if not candi_items:
                break  # 没有候选项集时退出循环
        else:
            break  # 没有频繁项集时退出循环
    return freq_items_list


def print_confident(items: set, confidence: float, data: list):
    """
    输出满足致信度的项
    :param items: 测试的项集
    :param confidence: 致信度
    :param data: 事物数据
    :return:
    """
    if len(items) > 1:
        for i in items:
            s2_set = set()
            for j in items:
                if j != i:
                    s2_set.add(j)
            s2 = get_support(s2_set, data)
            s1 = get_support(items, data)
            p = round(s1 / s2, 3)
            if p >= confidence:
                print('{} => {}的置信度为：{}'.format(s2_set, i, p))


if __name__ == '__main__':
    trans_data = deal_data(read_data(file_path))  # 从文件中读取事物数据
    print('找出数据集中支持度为{}的频繁项集，并从中找出置信度为{}以上的强关联关系'.format(support_num, confidence))
    freq_items_list = get_all_freq_items(trans_data, support_num)
    item_count = 1
    for items in freq_items_list:
        print('支持度为{}的频繁{}项集：'.format(support_num, item_count), end='')
        for i in items:
            print(i, end='  ')
        print()
        item_count += 1
    print('强关联关系：')
    for items in freq_items_list[1::]:
        for i in items:
            print_confident(i, confidence, trans_data)
